package com.one.blocks.trace.interceptor;

import com.one.blocks.common.util.SnowflakeIdGen;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import static com.one.blocks.trace.constant.TraceConstant.ID_KEY;

/**
 * @author linfuxin Created on 2022-02-14 10:47:10
 */
@Slf4j
public class TraceInterceptor extends HandlerInterceptorAdapter {

    private final SnowflakeIdGen traceIdGenerator;

    public TraceInterceptor(SnowflakeIdGen traceIdGenerator) {
        this.traceIdGenerator = traceIdGenerator;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod && needGenerate(request)) {
            long traceId = traceIdGenerator.nextId();
            request.setAttribute(ID_KEY, traceId);
            MDC.put(ID_KEY, Long.toString(traceId));
        }
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        if (handler instanceof HandlerMethod) {
            MDC.remove(ID_KEY);
            request.removeAttribute(ID_KEY);
        }
    }

    private boolean needGenerate(HttpServletRequest request) {
        // 请求转发不需要重新生成traceId
        Object traceKey;
        return (traceKey = request.getAttribute(ID_KEY)) == null || !traceKey.equals(MDC.get(ID_KEY));
    }
}